import cv2
import os
from ultralytics import YOLO

def detect_person_in_video(model, video_path, output_dir, skip):
    # 创建输出目录（如果不存在的话）
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 提取文件路径中的目录和文件名
    directory, filename = os.path.split(video_path)  # 分离出目录和文件名
    name_without_extension = os.path.splitext(filename)[0]  # 去掉扩展名（.mp4）

    # 从路径中提取上一级目录的名称
    parent_dir = os.path.basename(os.path.dirname(video_path))

    # 组合新的文件名：G:/vlog/outputs/parent_dir/name_without_extension.jpg
    output_filename = os.path.join(output_dir, f"{parent_dir}/{name_without_extension}.jpg")

    # 创建子目录（如果不存在的话）
    output_subdir = os.path.dirname(output_filename)
    if not os.path.exists(output_subdir):
        os.makedirs(output_subdir)

    cap = cv2.VideoCapture(video_path)

    frame_count = 0  # 帧计数器
    detection_count = 0  # 检测到人的次数
    detected = False  # 是否检测到人

    while True:
        ret, frame = cap.read()

        if not ret:
            break

        # 每隔 skip 帧进行一次推理
        if frame_count % skip == 0:
            frame_resized = cv2.resize(frame, (640, 360))
            res = model(frame_resized, verbose=False)  # 设置verbose=False避免打印日志

            cls = res[0].boxes.cls  # 检测到的物体类型
            xyxy = res[0].boxes.xyxy  # 检测到的物体位置

            detect = False  # 是否检测到人

            # 遍历所有的检测结果，判断是否有检测到人
            for i, t in enumerate(cls):
                if t == 0:  # 0 表示 person 类型
                    detect = True
                    break

            if detect:
                # 如果检测到人，保存图像并覆盖之前的图像
                cv2.imwrite(output_filename, frame)  # 保存当前帧图像
                detection_count += 1
                detected = True  # 标记检测到人

                # 如果已经检测到3次人，退出
                if detection_count >= 5:
                    break

        frame_count += 1  # 增加帧计数器

    cap.release()
    return detected

def list_mp4_files(model, skip, directory, output_dir):
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    output_file = os.path.join(output_dir, "stats.txt")
    print(output_file)
    with open(output_file, 'w', encoding='utf-8') as file:
        for root, dirs, files in os.walk(directory):
            for name in files:
                if name.lower().endswith('.mp4'):
                    video_path = os.path.join(root, name)
                    detect = detect_person_in_video(model, video_path, output_dir, skip)
                    if detect:
                        file.write(video_path + ' === person\n')
                        print(video_path, "=== detect person")
                    else:
                        file.write(video_path + 'none\n')
                        print(video_path, " none")

model = YOLO('yolov5su.pt')  # 加载YOLO模型
skip = 10

directory = r'E:/vlog/vlog-1208'
output_file = r'E:/vlog/vlog_1208_skip10'
list_mp4_files(model, skip, directory, output_file)
